package com.ccx.image.classification;

import android.annotation.SuppressLint;
import android.content.ContentResolver;
import android.database.Cursor;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.net.Uri;
import android.os.Build;
import android.provider.MediaStore;
import android.util.Base64;
import android.util.Log;
import android.util.Size;

import org.apache.cordova.CallbackContext;
import org.apache.cordova.CordovaArgs;
import org.apache.cordova.CordovaPlugin;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.io.ByteArrayOutputStream;
import java.util.HashMap;
import java.util.Map;

public class ImageClassification extends CordovaPlugin {

    public static final String TAG = "Cordova.Plugin.ImageClassification";

    protected Predictor predictor = new Predictor();

    @Override
    protected void pluginInitialize() {
        super.pluginInitialize();

        boolean res = onLoadModel();
        Log.i(TAG, "load model: " + res);

        Log.d(TAG, "plugin initialized.");
    }

    public boolean onLoadModel() {
        try{
            boolean success = predictor.init(cordova.getContext(), "models/mobilenet_v1_for_cpu", "labels/synset_words.txt", 1,
                    "LITE_POWER_HIGH",
                    "RGB",
                    new long[]{1,3,224,224}, new float[]{0.485f,0.456f,0.406f},
                    new float[]{0.229f,0.224f,0.225f});
            return success;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    @Override
    public boolean execute(String action, CordovaArgs args, CallbackContext callbackContext) throws JSONException {
        Log.d(TAG, String.format("%s is called. Callback ID: %s.", action, callbackContext.getCallbackId()));

        if (action.equals("classify")) {
            return classify(args, callbackContext);
        }
        if (action.equals("classifyAll")) {
            return classifyAll(args, callbackContext);
        }

        return false;
    }

    private boolean classify(CordovaArgs args, final CallbackContext callbackContext) {
        cordova.getThreadPool().execute(() -> {
            try {
                String imgBase64 = args.getString(0);
                imgBase64 = imgBase64.substring(imgBase64.indexOf(",") + 1);
                byte[] imgByte = Base64.decode(imgBase64, Base64.DEFAULT);
                Bitmap img = BitmapFactory.decodeByteArray(imgByte, 0, imgByte.length);
                predictor.setInputImage(img);
                boolean res = predictor.runModel();
                if(res){
                    callbackContext.success(String.format("%s,%s,%s", predictor.top1Result(), predictor.top2Result(), predictor.top3Result()));
                } else {
                    callbackContext.error("failed!");
                }
            } catch (Exception e) {
                callbackContext.error(e.getMessage());
            }
        });
        return true;
    }

    private boolean classifyAll(CordovaArgs args, final CallbackContext callbackContext) {
        cordova.getThreadPool().execute(() -> {
            try {
                ContentResolver cr = cordova.getContext().getContentResolver();
                if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.O) {
                    @SuppressLint("Recycle") Cursor cursor = cr.query(MediaStore.Images.Media.EXTERNAL_CONTENT_URI, new String[]{
                            MediaStore.Images.ImageColumns._ID,
                            MediaStore.Images.ImageColumns.DATA,
                            MediaStore.Images.ImageColumns.BUCKET_ID,
                            MediaStore.Images.ImageColumns.BUCKET_DISPLAY_NAME
                    }, null, null, MediaStore.Images.ImageColumns.DATE_MODIFIED + " DESC");
                    if (cursor == null) {
                        callbackContext.error("出错了");
                        return;
                    }
                    JSONArray images = new JSONArray();
                    Map<String, JSONArray> map = new HashMap<>();
                    for (String wordLabel : predictor.wordLabels) {
                        map.put(wordLabel.trim(), new JSONArray());
                    }
                    map.put("其他", new JSONArray());
                    if (cursor.moveToFirst()) {
                        do {
                            long imgId = cursor.getLong(0);
                            String imgPath = cursor.getString(1);
                            Bitmap bitmap = null;
                            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q) {
                                Uri uri = Uri.withAppendedPath(MediaStore.Images.Media.EXTERNAL_CONTENT_URI, Long.toString(imgId));
                                bitmap = cordova.getContext().getContentResolver().loadThumbnail(uri, new Size(224, 224), null);
                            } else {
                                bitmap = MediaStore.Images.Thumbnails.getThumbnail(cordova.getContext().getContentResolver(), imgId, MediaStore.Images.Thumbnails.MINI_KIND, null);
                            }
                            if (bitmap == null) {
                                continue;
                            }
                            predictor.setInputImage(bitmap);
                            boolean res = predictor.runModel();
                            if(res){
                                ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
                                bitmap.compress(Bitmap.CompressFormat.JPEG, 70, byteArrayOutputStream);
                                byte[] byteArray = byteArrayOutputStream.toByteArray();
                                String encoded = Base64.encodeToString(byteArray, Base64.DEFAULT);
                                String top1Result = predictor.top1Result();
                                String[] arr = top1Result.split("\\s+");
                                String category = arr[1];
                                float p = Float.parseFloat(arr[3]);
                                String imgBase64 = String.format("data:image/jpeg;base64,%s", encoded);
                                JSONObject obj = new JSONObject();
                                obj.put("imgId", imgId);
                                obj.put("imgPath", imgPath);
                                obj.put("imgBase64", imgBase64);
                                if(p < 0.8){
                                    map.get("其他").put(obj);
                                } else {
                                    map.get(category).put(obj);
                                }
                            }
                        } while (cursor.moveToNext());
                        for (String wordLabel : predictor.wordLabels) {
                            JSONObject obj = new JSONObject();
                            obj.put("name", wordLabel.trim());
                            obj.put("data", map.get(wordLabel.trim()));
                            images.put(obj);
                        }
                        JSONObject obj = new JSONObject();
                        obj.put("name", "其他");
                        obj.put("data", map.get("其他"));
                        images.put(obj);
                        callbackContext.success(images);
                    }
                    cursor.close();
                }
            } catch (Exception e) {
                callbackContext.error(e.getMessage());
            }
        });
        return true;
    }
}
